from __future__ import division, print_function

import os
import os.path as osp

import torch
from torch import nn
from torch.autograd import Variable
from torch.nn import Tensor

from retinaNet import retinaNet
from focalLoss import focalLoss

# define wrapper for retinaNet

class retinaWrapper(object):
    def __init__(self):
        super(retinaWrapper, self).__init__()

        self.is_cuda    = False

    def _init_weight(self):
        """
        initialize the weight for current network
        """

    def _validate(self):
        """
        iterate on the validation data to complete training
        generally, the validation data is the same as training data
        """

    def _val_iterate(self, batch_data):
        """
        iterate on current batch, forward,calculate loss, but not backward
        batch_data: the data of current batch
        """
        _loss   = self._iterate(batch_data, volatile=True)

    def _train(self, begin_epoch=0, end_epoch=None, savePt=False):
        """
        iterate on the train data to complete training
        """
        # initialize the training data
        self._init_train_data()
        # training mode
        self._det_net.train()
        if end_epoch is None:
            end_epoch   = self._learnCon.epoch_num
        print("begin training")

        #TODO, step the learning rate
        for i in range(begin_epoch, end_epoch):
            print("{}-th epoch start".format(i))
            for j, j_batch_data in enumerate(self._det_loader):
                j_loss_val  = self._train_iterate(j_batch_data)
                print("  epoch: {} --- batch: {}  current loss is: {}".format(i, j, j_loss_val))
            # TODO, mean value for multiple losses
            
            # save temporary model parameters
            if (i+1) % self._learnCon.saveEpochGap  == 0:
                # save model
                if savePt:
                    modelDir    = self._learnCon.saveModelDir
                    modelName   = "epoch-{:03d}.format(i)"
                    self.saveModel(modelDir, modelName)

    def _train_iterate(self, batch_data):
        """
        iterate on current batch, forward,calculate loss, then backward
        batch_data: the data of current batch
        """
        self._opt.zero_grad()
        _loss   = self._iterate(batch_data, volatile=True)
        _loss.backward()
        self._opt.step()

        # TODO, detailed loss
        _loss_val   = _loss.cpu().data.numpy()[0]
        return _loss_val

    def _iterate(self, batch_data, volatile=False):
        """
        iterate a batch, forward and calculate loss
        batch_data: the data of current batch
        """
        _img, loc_target, cls_target    = batch_data[0:3]
        _img_var, _loc_target_var, _cls_target_var  = \
            Variable(_img, volatile=volatile),\
            Variable(_loc_target_var, volatile=volatile),\
            Variable(_cls_target_var, volatile=volatile)
        if self.is_cuda:
            _img_var, _loc_target_var, _cls_target_var  = \
            _img_var.cuda(), _loc_target_var.cuda(), _cls_target_var.cuda()

        _loc_pred_var, _cls_pred_var    = self._det_net(_img_var)
        _loss   = self._det_loss(_loc_pred_var, _loc_target_var, _cls_pred_var, _cls_target_var)
        return _loss

    def loadModel(self, modelDir, modelName="retinaNet"):
        detNetPath  = osp.join(modelDir, modelName+".pt")
        tmp         = torch.load(detNetPath)
        orin_cuda   = self.is_cuda
        self._det_net.cpu()
        self._det_net.load_state_dict(tmp)
        if orin_cuda:
            self._det_net.cuda()

    def saveModel(self, modelDir, modelName="retinaNet"):
        """only save _det_net"""
        if not osp.exists(modelDir):
            os.makedirs(modelDir)
            print("create modelDir: {} for saving retinaNet".format(modelDir))
        detNetPath  = osp.join(modelDir, modelName+".pt")
        orin_cuda   = self.is_cuda
        self._det_net.cpu()
        torch.save(self._det_net.state_dict(), detNetPath)
        if orin_cuda:
            self._det_net.cuda()
        

    def cuda(self):
        """
        assert the computation and variable will be on cpu
        """
        self.is_cuda    = True
        self._det_net.cuda()

    def cpu(self):
        """
        assert the computation and variable will be on cpu
        """
        self.is_cuda    = False
        self._det_net.cpu()

    def _init_hyperparams(self):
        """
        initialize the hyperparams for training current network, etc
        """
        self._learnCon  = config.train.learn

    def _init_train_data(self):
        """
        initialize dataset for training
        """
        self._det_data  =
        self._det_loader=

    def _init_net(self):
        """
        initialize the network structure for current network
        """
        self._det_net   = retinaNet()

    def _init_loss(self):
        """
        initialize the loss configuration for current network
        """
        self._det_loss  = focalLoss()

    def _init_opt(self):
        """
        initialize the optimizer for current network
        """
        self._opt       = nn.SGD(self._det_net.parameters(), lr=self._learnCon.lr, momentum=self._learnCon.momentum, weight_decay=self._learnCon.weight_decay)
